package grpc

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
	"zhejianglab/gin-apiserver/model"
	"zhejianglab/gin-apiserver/pkg/logger"
	pb "zhejianglab/gin-apiserver/proto"

	"github.com/tal-tech/go-zero/core/collection"
)

type RpcService struct{}

var SuperResolutionJobQueue = collection.NewQueue(128)
var ObjectDetectionJobQueue = collection.NewQueue(128)
var ObjectTrackJobQueue = collection.NewQueue(128)
var ReIDJobQueue = collection.NewQueue(128)
var ActionRecognitionJobQueue = collection.NewQueue(128)

var SuperResolutionJobChan chan pb.Message
var ObjectDetectionJobChan chan pb.Message
var ObjectTrackJobChan chan pb.Message
var ReIDJobChan chan pb.Message
var ActionRecognitionJobChan chan pb.Message
var AllJobChan chan pb.Message

func init() {
	SuperResolutionJobChan = make(chan pb.Message, 50)
	ObjectDetectionJobChan = make(chan pb.Message, 50)
	ObjectTrackJobChan = make(chan pb.Message, 50)
	ReIDJobChan = make(chan pb.Message, 50)
	ActionRecognitionJobChan = make(chan pb.Message, 50)
	AllJobChan = make(chan pb.Message, 50)
}
func closeChan() {
	close(SuperResolutionJobChan)
	close(ObjectTrackJobChan)
	close(ObjectTrackJobChan)
	close(ReIDJobChan)
	close(ActionRecognitionJobChan)
	close(AllJobChan)
	fmt.Println("close channel")
}

var wg sync.WaitGroup

func (s *RpcService) Start(srv pb.ServiceTask_StartServer) error {
	logger.Log.Info("start new server")
	wg.Add(1)
	ctx, cancel := context.WithCancel(srv.Context())
	go receiveTask(ctx, srv, &wg)
	wg.Wait()
	fmt.Println("server end cancel goroutine........")
	cancel()
	// closeChan()
	return nil
}

func receiveTask(ctx context.Context, srv pb.ServiceTask_StartServer, wg *sync.WaitGroup) {
	defer wg.Done()
	for {

		req, err := srv.Recv()
		if err == io.EOF {
			fmt.Println("recv EOF \n")
		}
		if err != nil {
			fmt.Println("recv \n", err.Error())
			return
		}
		if req != nil {
			// 心跳
			if req.MessageType == pb.MesageType_MESSAGETYPE_HEARTBEAT {
				fmt.Printf("RequestId:%d heartbeat！\n", req.RequestId)
			}
			// 注册
			if req.MessageType == pb.MesageType_MESSAGETYPE_REGISTRY {
				switch req.JobType {
				case 0:
					logger.Log.Info("注册图像超分辨任务！")
					go sendTask(ctx, srv, SuperResolutionJobChan)
				case 1:
					logger.Log.Info("注册目标检测任务")
					go sendTask(ctx, srv, ObjectDetectionJobChan)
				case 2:
					logger.Log.Info("注册目标跟踪任务")
					go sendTask(ctx, srv, ObjectTrackJobChan)
				case 3:
					logger.Log.Info("注册ReID任务")
					go sendTask(ctx, srv, ReIDJobChan)
				case 4:
					logger.Log.Info("注册行为分析任务")
					go sendTask(ctx, srv, ActionRecognitionJobChan)
				case 5:
					logger.Log.Info("注册算法任务")
					go sendTask(ctx, srv, AllJobChan)
				}

			}
			// 上报任务状态
			if req.MessageType == pb.MesageType_MESSAGETYPE_REPORT_JOB {
				fmt.Printf("RequestId:%d 上报任务 %d 进度 %d！\n", req.RequestId, req.AlgorithmProcess.JobId, req.AlgorithmProcess.Process)
				jobProcess := model.JobInfo{
					Progress: int64(req.AlgorithmProcess.Process),
				}
				model.UpdateJob(int(req.AlgorithmProcess.JobId), jobProcess)
			}
			// 上报任务结果
			if req.MessageType == pb.MesageType_MESSAGETYPE_REPORT_RESULT {
				fmt.Printf("RequestId:%d 上报job %d结果！\n", req.RequestId, req.AlgorithmResult.JobId)
				jobResult := model.JobInfo{
					ResultPath:       req.AlgorithmResult.ResultPath,
					ResultStatistics: req.AlgorithmResult.ResultStatistics,
				}
				// 超分辨算法没有进度信息，完成进度就是100
				if req.JobType == pb.JobType_JOB_SUPERRESOLUTION {
					jobResult.Progress = 100
				}
				model.UpdateJob(int(req.AlgorithmResult.JobId), jobResult)
				// 将算法任务结果写到数据库中
			}
			// 任务出错，进度写入-1，返回错误原因
			if req.ErrorCode != pb.ErrorCode_ERRORCODE_SUCCESS {
				fmt.Printf("Error %s", req.ErrorStr)
				jobResult := model.JobInfo{
					Progress:  -1,
					ErrorCode: int64(req.ErrorCode),
					ErrorMsg:  req.ErrorStr,
				}
				model.UpdateJob(int(req.AlgorithmResult.JobId), jobResult)
			}
			fmt.Printf("receive task from client,the job type is:%d \n", req.JobType)
		}
	}

}
func sendTask(ctx context.Context, srv pb.ServiceTask_StartServer, ch chan pb.Message) {
jobLoop:
	for {
		select {
		case job := <-ch:
			fmt.Println("task to handle")
			job.RequestId = 100
			err := srv.Send(&job)
			if err != nil {
				fmt.Println("grpc err", err)
				SuperResolutionJobChan <- job
				return
			}
			logger.Log.Info(fmt.Sprintf("发送%d处理任务！", job.JobType))
		//case job := <-ObjectDetectionJobChan:
		//	fmt.Println("task to handle")
		//	job.RequestId = 100
		//	err := srv.Send(&job)
		//	if err != nil {
		//		fmt.Println("grpc err", err)
		//		ObjectDetectionJobChan <- job
		//		return
		//	}
		//	logger.Log.Info(fmt.Sprintf("发送目标检测处理任务！"))
		//case job := <-ObjectTrackJobChan:
		//	fmt.Println("task to handle")
		//	job.RequestId = 100
		//	err := srv.Send(&job)
		//	if err != nil {
		//		fmt.Println("grpc err", err)
		//		ObjectTrackJobChan <- job
		//		return
		//	}
		//	logger.Log.Info(fmt.Sprintf("发送目标跟踪处理任务！"))
		//case job := <-ReIDJobChan:
		//	fmt.Println("task to handle")
		//	job.RequestId = 100
		//	err := srv.Send(&job)
		//	if err != nil {
		//		fmt.Println("grpc err", err)
		//		ReIDJobChan <- job
		//		return
		//	}
		//	logger.Log.Info(fmt.Sprintf("发送ReID处理任务"))
		//case job := <-ActionRecognitionJobChan:
		//	fmt.Println("task to handle")
		//	job.RequestId = 100
		//	err := srv.Send(&job)
		//	if err != nil {
		//		fmt.Println("grpc err", err)
		//		ActionRecognitionJobChan <- job
		//		return
		//	}
		//	logger.Log.Info(fmt.Sprintf("发送行为分析处理任务"))
		case <-time.After(time.Millisecond * 10):
		case <-ctx.Done():
			fmt.Println(ctx.Err()) // prints "context deadline exceeded"
			break jobLoop
		}
	}

}

func handleTask(parent context.Context, srv pb.ServiceTask_StartServer, jobType pb.JobType) {
	wokeQueue := &collection.Queue{}
	switch jobType {
	case 0:
		logger.Log.Info("注册图像超分辨任务！")
		wokeQueue = SuperResolutionJobQueue
	case 1:
		logger.Log.Info("注册目标检测任务")
		wokeQueue = ObjectDetectionJobQueue
	case 2:
		logger.Log.Info("注册目标跟踪任务")
		wokeQueue = ObjectTrackJobQueue
	case 3:
		logger.Log.Info("注册ReID任务")
		wokeQueue = ReIDJobQueue
	case 4:
		logger.Log.Info("注册行为分析任务")
		wokeQueue = ActionRecognitionJobQueue
	}
loop:
	for {
		select {
		case <-parent.Done():
			fmt.Println("cancelled by sub goroutine")
			break loop
		default:
			JResult, ok := wokeQueue.Take()
			if ok {
				fmt.Println("handleTask")
				resAcc, ok := JResult.(pb.Message)
				if ok {
					resAcc.RequestId = 100
					err := srv.Send(&resAcc)
					if err != nil {
						fmt.Println("send grpc task error", err)
					}
					logger.Log.Info(fmt.Sprintf("发送处理任务%d！", resAcc.AlgorithmJob.JobId))
				} else {
					err := srv.Send(&pb.Message{
						RequestType: 0,
						ErrorCode:   pb.ErrorCode_ERRORCODE_RUN_FAILED,
					})
					if err != nil {
						fmt.Println("Assert failed！")
					}

				}
			}
		}

	}
	fmt.Println(" sub goroutine finished")
}

func handleReIDTask(req *pb.Message, srv pb.ServiceTask_StartServer) error {
	JResult, ok := ReIDJobQueue.Take()
	if ok {
		resAcc, ok := JResult.(pb.Message)
		if ok {
			resAcc.RequestId = req.RequestId
			resAcc.RequestType = pb.RequestType_REQUESTTYPE_ACK
			err := srv.Send(&resAcc)
			if err != nil {
				return err
			}
			logger.Log.Info(fmt.Sprintf("发送ReID处理任务%d！", resAcc.AlgorithmJob.JobId))
		} else {
			err := srv.Send(&pb.Message{
				RequestType: 0,
				ErrorCode:   pb.ErrorCode_ERRORCODE_RUN_FAILED,
			})
			if err != nil {
				return err
			}
			return errors.New("Assert failed！")
		}
	} else {
		err := srv.Send(&pb.Message{
			RequestType: 0,
			ErrorCode:   pb.ErrorCode_ERRORCODE_RUN_FAILED,
		})
		if err != nil {
			return err
		}
		return errors.New("Take element from ReID queue failed！")
	}
	return nil
}
